import subprocess
import os
import sys
import json


class AtcOm():
    def __init__(self):
        self.op_debug_level = 0
        self.op_precision_mode = None
        self.atc_log = "debug"

    def save_json(self, data, file_path):
        with open(file_path, "w") as f:
            json.dump(data, f, indent=4)

    def run_atc(self, params):
        format_str = "--%s=%s"
        param_list = []
        for key, value in params.items():
            param_str = format_str % (key, value)
            param_list.append(param_str)

        atc_param = " ".join(param_list)
        atc_cmd = "atc %s" % atc_param
        return_code = subprocess.call(atc_cmd, shell=True, stderr=subprocess.PIPE, timeout=60000)
        if return_code != 0:
            print("\033[31m[> 错误 <]\033[0m" + "编译算子失败")
            print("\033[31m[> 错误 <]\033[0m" + "编译算子命令:", atc_cmd)
            print("\033[31m[> 错误 <]\033[0m" + "相关atc工具命令,请翻阅手册查看")
            print("\033[31m[> 错误 <]\033[0m" + "1.重点检查自己写的脚本能不能与算子的IR对应上.")
            print("\033[31m[> 错误 <]\033[0m" + "2.日志排查:source set_log.sh")
            sys.exit(1)
        return return_code

    def generate_json(self, params, op_dir, compile_dict=None):
        json_name = params["op"].lower() + ".json"
        json_path = os.path.join(op_dir, json_name)
        if compile_dict is not None:
            self.save_json([compile_dict, params], json_path)
        else:
            self.save_json([params], json_path)
        return json_path

    def generate_om(self, params, op_dir, compile_flag, soc_version="Ascend910", deterministic=0, op_debug_config=""):
        if deterministic == 1:
            print("\033[93m[> 检查 <]\033[0m", "已开启确定性计算, tiktools.deterministic = 1")
        compile_dict = None
        if self.op_debug_level == 0:
            compile_flag = 1
        if compile_flag:
            compile_dict = {"compile_flag": compile_flag}
        json_path = self.generate_json(params, op_dir, compile_dict)
        om_params = {"singleop": json_path, "output": op_dir, "soc_version": soc_version,
                     "log": self.atc_log, "op_debug_level": self.op_debug_level,
                     "precision_mode": "allow_fp32_to_fp16", "deterministic": deterministic}
        if self.op_precision_mode is not None:
            om_params["op_precision_mode"] = self.op_precision_mode
        if op_debug_config:
            om_params["op_debug_config"] = op_debug_config

        return_code = self.run_atc(om_params)
        return return_code


atc = AtcOm()
